import numpy as np
import sys
import cv2
import tritonclient.http as httpclient
import base64
import random
import os

categories = ["person", "cat", "dog", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
              "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
              "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
              "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
              "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
              "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
              "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
              "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
              "hair drier", "toothbrush"]


def get_box(triton_client):

    # file_name = "dog.jpg"
    #存放图片的路径
    #dirname = "person_cat_dog_images_valid-1"
    dirname = "/home/huangzhe/QuestCore_v0.3.4/Release/rpc_HeadChefhatMouse_tmp/images/test/test"
    img_dir = os.listdir(dirname)
    for img in img_dir:
        # file_name = ["person_cat_dog_images_valid/000000000139.jpg"]
        file_name = [dirname+"/"+img]
        input_base64_matrix = []
        img = cv2.imread(file_name[0])
        # input_base64_matrix.append([x for x in base64.b64encode(open("dog.jpg","rb").read())])
        input_base64_matrix.append([x for x in base64.b64encode(open(file_name[0],"rb").read())])
        # input_base64_matrix.append([x for x in base64.b64encode(open(file_name[1],"rb").read())])
        input_base64_matrix_len_list = [len(l) for l in input_base64_matrix]
        max_len = max(input_base64_matrix_len_list)
        input_base64_matrix = list(map(lambda l:l + [0]*(max_len - len(l)), input_base64_matrix))
        input_base64_matrix = np.array(input_base64_matrix,dtype=np.uint8)
        print(input_base64_matrix.shape)

        # 576,768 （h,w)
        input_param = [(img.shape[0],img.shape[1],x) for x in input_base64_matrix_len_list]
        # input_param = [(360,640,x) for x in input_base64_matrix_len_list]
        input_param = np.array(input_param, dtype=np.uint32)
        print(input_param)

        input_size = np.array([[352,640] for _ in range(len(file_name))],dtype=np.uint16)

        inputs =[]
        inputs.append(httpclient.InferInput('input_base64', list(input_base64_matrix.shape), "UINT8"))
        inputs.append(httpclient.InferInput('input_param', list(input_param.shape), "UINT32"))
        inputs.append(httpclient.InferInput('input_size', list(input_size.shape), "UINT16"))
        inputs[0].set_data_from_numpy(input_base64_matrix)
        inputs[1].set_data_from_numpy(input_param)
        inputs[2].set_data_from_numpy(input_size)

        outputs = []
        outputs.append(httpclient.InferRequestedOutput('output_yolo_res_nms_unresized'))
        results = triton_client.infer(
            model_name="ensemble", inputs=inputs,  outputs=outputs, headers={"test": "1"}
        )
        output_yolo_res_nms_unresized = results.as_numpy("output_yolo_res_nms_unresized")
        # print("output_yolo_res_nms_unresized.shape: ",output_yolo_res_nms_unresized.shape)
        count = int(output_yolo_res_nms_unresized[0][0])
        # for i in range(len(output_yolo_res_nms_unresized)):
        #     with open('test.txt', 'a') as f:
        #         f.write(str(file_name))
        #         f.write('\t')
        #     for j in range(len(output_yolo_res_nms_unresized[0])):
        #         with open('test.txt', 'a') as f:
        #             f.write(str(output_yolo_res_nms_unresized[i][j]))
        #         if j == 6*int(count):
        #             with open('test.txt', 'a') as f:
        #                 f.write("\n")
        #             break
        #     break
        curitem = output_yolo_res_nms_unresized[0]
        for i in range(int(curitem[0])):
            outstr = ""
            outstr += file_name[0]
            curobj = curitem[6*i+1:6*i+7]
            # print("curobj:",curobj)
            clas = int(curobj[5])
            curbox = curobj[0:4]
            y = [0 for _ in range(4)]
            y[0] = 0 if (curbox[0] - curbox[2] / 2) < 0 else (curbox[0] - curbox[2] / 2)
            y[2] = img.shape[1] if (curbox[0] + curbox[2] / 2) > img.shape[1] else (curbox[0] + curbox[2] / 2)
            y[1] = 0 if (curbox[1] - curbox[3] / 2) <0 else (curbox[1] - curbox[3] / 2)
            y[3] = img.shape[0] if (curbox[1] + curbox[3] / 2) > img.shape[0] else curbox[1] + curbox[3] / 2
            newbox = [y[0], y[1], y[2], y[3]]
            conf = curobj[4]
            outstr += "\t"
            outstr += str(clas)
            outstr += "\t"
            outstr += str(conf)
            outstr += "\t"
            outstr += str(newbox)
            outstr += "\n"
            with open('test.txt', 'a') as f:
                f.write(outstr)



        for fileindex in range(len(file_name)):
            print("boxes num:", output_yolo_res_nms_unresized[fileindex][0])
            boxes = np.reshape(output_yolo_res_nms_unresized[fileindex][1:], (-1, 6))
            print(boxes.shape)
            result_boxes = boxes[:,:4]
            result_scores = boxes[:,4]
            result_classid = boxes[:,5]
            print(result_scores[:10])

            image_raw = cv2.imread(file_name[fileindex])
            for index,box in enumerate(result_boxes):
                y = [0 for _ in range(4)]
                y[0] = box[0] - box[2]/2
                y[2] = box[0] + box[2]/2
                y[1] = box[1] - box[3]/2
                y[3] = box[1] + box[3]/2
                # print(y)
                plot_one_box(
                    y,
                    image_raw,
                    label="{}:{:.2f}".format(
                        categories[int(result_classid[index])], result_scores[index]
                    ),
                )
            print("index",fileindex)
            cv2.imwrite('result/'+file_name[fileindex].split('/')[-1],image_raw)
            print('result/'+file_name[fileindex].split('/')[-1])

def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param:
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return
    """
    tl = (
        line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
            )

if __name__ == '__main__':
    # triton_client = httpclient.InferenceServerClient(url="127.0.0.1:42000", verbose=False) #工作站上run，请求工作站服务
    triton_client = httpclient.InferenceServerClient(url="10.143.165.45:31004", verbose=False) #工作站上run，请求工作站服务
    get_box(triton_client)
